import akshare as ak
import pandas as pd
from datetime import datetime, timedelta

def get_stock_data(symbol, days=60):
    """获取股票历史数据"""
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days)
    
    # 格式化日期为YYYYMMDD格式
    start_date_str = start_date.strftime('%Y%m%d')
    end_date_str = end_date.strftime('%Y%m%d')
    
    # 提取纯股票代码（去掉交易所后缀）
    pure_symbol = symbol.split('.')[0]
    print(f"尝试使用AkShare获取股票数据: {pure_symbol}, 时间范围: {start_date_str} 到 {end_date_str}")
    
    try:
        # 使用AkShare获取A股历史数据
        # adjust="qfq"表示前复权
        data = ak.stock_zh_a_hist(symbol=pure_symbol, period="daily", start_date=start_date_str, end_date=end_date_str, adjust="qfq")
        
        if data.empty:
            print(f"股票 {pure_symbol} 数据为空")
            return None
        
        # 重命名列以匹配原有的yfinance输出格式
        data = data.rename(columns={
            '日期': 'Date',
            '开盘': 'Open',
            '收盘': 'Close',
            '最高': 'High',
            '最低': 'Low',
            '成交量': 'Volume',
            '成交额': 'Turnover'
        })
        
        # 将日期列设置为索引
        data['Date'] = pd.to_datetime(data['Date'])
        data.set_index('Date', inplace=True)
        
        print(f"成功获取股票 {pure_symbol} 数据，数据长度: {len(data)}")
        return data
    except Exception as e:
        error_msg = str(e)
        print(f"获取股票 {pure_symbol} 数据错误: {error_msg}")
        # 抛出异常以便上层处理
        raise Exception(f"AkShare获取数据失败: {error_msg}")


def process_stock_symbol(symbol):
    """处理股票代码格式"""
    # 如果已经包含交易所后缀，直接返回
    if symbol.endswith('.SS') or symbol.endswith('.SZ'):
        return symbol
    
    # 根据股票代码长度判断交易所（简单判断）
    # 6开头的通常是上交所，0开头或3开头的通常是深交所
    if len(symbol) == 6:
        if symbol.startswith('6'):
            return symbol + '.SS'  # 上交所
        elif symbol.startswith('0') or symbol.startswith('3'):
            return symbol + '.SZ'  # 深交所
    
    return symbol
